// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <pollux/common/utils.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/vector/vector_fuzzer.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions;

namespace {

using fuzzer::UTF8CharList;

class StringAsciiUTFFunctionBenchmark
    : public functions::test::FunctionBenchmarkBase {
 public:
  StringAsciiUTFFunctionBenchmark() : FunctionBenchmarkBase() {
    functions::prestosql::registerStringFunctions();
  }

  void runUpperLower(const std::string& fnName, bool utf) {
    melon::BenchmarkSuspender suspender;

    VectorFuzzer::Options opts;
    if (utf) {
      opts.charEncodings.clear();
      opts.charEncodings = {UTF8CharList::UNICODE_CASE_SENSITIVE};
    }

    opts.stringLength = 100;
    opts.vectorSize = 100'000;
    VectorFuzzer fuzzer(opts, execCtx_.pool());
    auto vector = fuzzer.fuzzFlat(VARCHAR());

    auto row_vector = vectorMaker_.row_vector({vector});
    auto exprSet =
        compileExpression(fmt::format("{}(c0)", fnName), row_vector->type());

    suspender.dismiss();

    doRun(exprSet, row_vector);
  }

  void runSubStr(bool utf) {
    melon::BenchmarkSuspender suspender;

    VectorFuzzer::Options opts;
    if (utf) {
      opts.charEncodings.clear();
      opts.charEncodings = {
          UTF8CharList::UNICODE_CASE_SENSITIVE,
          UTF8CharList::EXTENDED_UNICODE,
          UTF8CharList::MATHEMATICAL_SYMBOLS};
    }

    opts.stringLength = 100;
    opts.vectorSize = 10'000;
    VectorFuzzer fuzzer(opts, execCtx_.pool());
    auto vector = fuzzer.fuzzFlat(VARCHAR());

    auto positionVector = BaseVector::create_constant(
        INTEGER(), 25, opts.vectorSize, execCtx_.pool());

    auto row_vector = vectorMaker_.row_vector({vector, positionVector});

    auto exprSet = compileExpression("substr(c0, c1)", row_vector->type());

    suspender.dismiss();
    doRun(exprSet, row_vector);
  }

  void runLPadRPad(const std::string& fnName, bool utf) {
    melon::BenchmarkSuspender suspender;

    VectorFuzzer::Options opts;
    if (utf) {
      opts.charEncodings.clear();
      opts.charEncodings = {
          UTF8CharList::UNICODE_CASE_SENSITIVE,
          UTF8CharList::EXTENDED_UNICODE,
          UTF8CharList::MATHEMATICAL_SYMBOLS};
    }

    opts.stringLength = 10;
    opts.vectorSize = 10'000;
    VectorFuzzer fuzzer(opts, execCtx_.pool());
    auto stringVector = fuzzer.fuzzFlat(VARCHAR());
    auto padStringVector = fuzzer.fuzzFlat(VARCHAR());

    auto sizeVector = BaseVector::create_constant(
        INTEGER(), 55, opts.vectorSize, execCtx_.pool());

    auto row_vector =
        vectorMaker_.row_vector({stringVector, sizeVector, padStringVector});

    auto exprSet = compileExpression(
        fmt::format("{}(c0, c1, c2)", fnName), row_vector->type());

    suspender.dismiss();
    doRun(exprSet, row_vector);
  }

  void doRun(ExprSet& exprSet, const RowVectorPtr& row_vector) {
    uint32_t cnt = 0;
    for (auto i = 0; i < 100; i++) {
      cnt += evaluate(exprSet, row_vector)->size();
    }
    melon::doNotOptimizeAway(cnt);
  }
};

BENCHMARK(utfLower) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runUpperLower("lower", true);
}

BENCHMARK_RELATIVE(asciiLower) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runUpperLower("lower", false);
}

BENCHMARK(utfUpper) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runUpperLower("upper", true);
}

BENCHMARK_RELATIVE(asciiUpper) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runUpperLower("upper", false);
}

BENCHMARK(utfSubStr) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runSubStr(true);
}

BENCHMARK_RELATIVE(asciiSubStr) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runSubStr(false);
}

BENCHMARK(utfLPad) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runLPadRPad("lpad", true);
}

BENCHMARK_RELATIVE(aciiLPad) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runLPadRPad("lpad", false);
}

BENCHMARK(utfRPad) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runLPadRPad("rpad", true);
}

BENCHMARK_RELATIVE(aciiRPad) {
  StringAsciiUTFFunctionBenchmark benchmark;
  benchmark.runLPadRPad("rpad", false);
}
} // namespace

// Preliminary release run, before ascii optimization.
//============================================================================
//../../pollux/functions/prestosql/benchmarks/StringAsciiUTFFunctionBenchmarks.cpprelative
// time/iter  iters/s
//============================================================================
// utfLower                                                    67.71ms    14.77
// asciiLower                                        99.84%    67.82ms    14.75
// utfUpper                                                    67.75ms    14.76
// asciiUpper                                        98.22%    68.98ms    14.50
//============================================================================
int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};

  melon::runBenchmarks();
  return 0;
}
